import numpy as np
from scipy.optimize import minimize

def w_(x,a,b,sigma,rho,m):
    return a + b*( rho*(x-m) + np.sqrt((x-m)**2+sigma**2))

def dw(x,a,b,sigma,rho,m):
    return b * (rho + (x-m)/np.sqrt((x-m)**2+sigma**2))

def dw2(x,a,b,sigma,rho,m):
    sqrt_part = np.sqrt((x-m)**2+sigma**2)
    return b * (1/sqrt_part-(x-m)**2/sqrt_part**3)

def durrleman_condition(x,a,b,sigma,rho,m,epsilon=0):
    cond = (1-x*dw(x,a,b,sigma,rho,m)/2/w_(x,a,b,sigma,rho,m))**2\
        - dw(x,a,b,sigma,rho,m)**2/4*(1/w_(x,a,b,sigma,rho,m)+1/4)\
        + dw2(x,a,b,sigma,rho,m)/2
    return np.sum(np.where(cond>=epsilon,0,cond-epsilon))

def durrleman_condition_list(x,a,b,sigma,rho,m):
    return (1-x*dw(x,a,b,sigma,rho,m)/2/w_(x,a,b,sigma,rho,m))**2\
        - dw(x,a,b,sigma,rho,m)**2/4*(1/w_(x,a,b,sigma,rho,m)+1/4)\
        + dw2(x,a,b,sigma,rho,m)/2

def raw2JW(a,b,sigma,rho,m,tau):
    v = (a+b*(-rho*m+np.sqrt(m**2+sigma**2)))/tau
    sw = np.sqrt(v * tau)
    ps = 1/sw * b/2 * (rho-m/np.sqrt(m**2+sigma**2))
    p = 1/sw * b *(1-rho)
    c = 1/sw * b *(1+rho)
    vh = 1/tau * (a+ b*sigma*np.sqrt(1-rho**2))
    return v,ps,p,c,vh

def JW2raw(v,ps,p,c,vh,tau):
    sw = np.sqrt(v * tau)
    b = sw/2*(c+p)
    rho = (c-p)/(c+p)
    beta = rho - 2*ps*sw/b
    alpha = np.sign(beta) * np.sqrt(1/beta**2-1)
    m = (v-vh) * tau/b/(-rho+np.sign(alpha)*np.sqrt(1+alpha**2)-alpha*np.sqrt(1-rho**2))
    if m ==0 :
        sigma = (v * tau - a)/b
    else:
        sigma = alpha * m
    a = vh * tau - b*sigma*np.sqrt(1-rho**2)
    return a,b,sigma,rho,m

def func_w_sqrt_penalty_raw(w_fit_old,x,v,ps,p,c,vh,tau,penalty = 1e2,epsilon=0):
    a,b,sigma,rho,m = JW2raw(v,ps,p,c,vh,tau)
    w_fit = w_(x,a,b,sigma,rho,m)
    cond = durrleman_condition(x,a,b,sigma,rho,m,epsilon=epsilon)
    return - penalty*cond + np.sum((w_fit-w_fit_old)**2) 

def calibration_adjust(x,a,b,sigma,rho,m,tau,penalty=1e2,method='TR',cond_tol=0,verbose = False):
    # Adjust Is Needed ?
    if penalty<=0 or durrleman_condition(x,a,b,sigma,rho,m,epsilon=cond_tol)>=0:
        if verbose:
            print('Adjust not need: stage 1 durr cond passed')
        return a,b,sigma,rho,m
    
    v,ps,p,c,vh = raw2JW(a,b,sigma,rho,m,tau)
    if verbose:
        print('JW params:',v,' ',ps,' ',p,' ',c,' ',vh)
    # Repick c & vhat
    c_new = p + 2 * ps
    vh_new = v * 4*p*c_new/(p+c_new)**2
    if verbose:
        print('old: ',c,' ',vh)
        print('repicked new: ',c_new,' ',vh_new)

    w_fit_old = w_(x,a,b,sigma,rho,m)
    func_w_sqrt_penalty = lambda arg: func_w_sqrt_penalty_raw(w_fit_old,x,v,ps,p,arg[0],arg[1],tau,penalty,epsilon=cond_tol)

    region = [[min(c_new,c),max(c_new,c)],[min(vh_new,vh),max(vh_new,vh)]]
    
    if method == 'Nelder-Mead':
        opt_out = minimize(func_w_sqrt_penalty,x0=[c,vh],method='Nelder-Mead',\
            bounds= region,\
            options={'fatol':1e-10,'xatol':1e-5})
        a,b,sigma,rho,m = JW2raw(v,ps,p,opt_out.x[0],opt_out.x[1],tau)
    elif method == 'TR':
        opt_out = minimize(func_w_sqrt_penalty,x0=[c,vh],method='trust-constr',\
            bounds= region,\
            options={'gtol':1e-10})
        a,b,sigma,rho,m = JW2raw(v,ps,p,opt_out.x[0],opt_out.x[1],tau)
    else:
        assert 1<0,"Calib adjust.py: Invalid Method"

    if verbose:
        print('adj_raw params :',a,' ',b,' ',sigma,' ',rho,' ',m)
    if durrleman_condition(x,a,b,sigma,rho,m,epsilon=cond_tol)<0:
        print('Warning: Arbitrage not eliminated after Adjust, durr cond failed')
    return a,b,sigma,rho,m